library(gplots)

params <- matrix(nrow=1080,ncol=6)
dpl <- c(0,0.25,0.5,0.75,1)
scl <- c(0.001,0.01,0.1,0.5,1,2)
sml <- c(10^-5,10^-6,10^-7,10^-8)
nml <- c(10^-4,10^-5,10^-7)
sw <- 0

for (divprob in 1:5){
  for (selcoeff in 1:6){
    for (selmu in 1:4){
      for (neutmu in 1:3){
        for (replicates in 1:3){
          sw <- sw + 1
          params[sw,1] <- sw
          params[sw,2] <- dpl[divprob]
          params[sw,3] <- scl[selcoeff]
          params[sw,4] <- sml[selmu]
          params[sw,5] <- nml[neutmu]
          params[sw,6] <- replicates
        }
      }
    }	
  }
}

res <- matrix(data=NA,nrow=120,ncol=36)
row <- 0
for (smui in 1:length(sml)){
  smu <- sml[smui]
  for (ss in 1:length(scl)){
    s <- scl[ss]
    for (r in 1:length(dpl)){
      column <- 0
      row <- row + 1
      for (nmu in 1:length(nml)){
        for (repl in 1:3){
          sw <- params[params[,3]==s&params[,4]==smu&params[,2]==dpl[r]&params[,5]==nml[nmu]&params[,6]==repl,1]
          # get the selective clones
          d <- read.csv(paste("../sim",sw,"/sim",sw,".frequencies",sep=""),header=FALSE)
          # get clone_ids and frequencies : end of simulation, frequency higher than 1%, selective clone (last hit is a selective locus)
          dat <- d[d[,2]==7300&d[,3]>0.01&d[,1]>0&d[,4]>104&d[,4]<110,c(1,3)]
      # get s from clone_ids, from .clones the number of beneficial hits that the clone has
          bh <- read.csv(paste("../sim",sw,"/sim",sw,".clones",sep=""),header=FALSE)
      # match subset to superset, get indices of the subset elements in the superset
          ind <- match(dat[,1],bh[,1])
          clones.s <- bh[ind,c(1,3)]

      # Get index of clones with maximum s -> last category, the most hits
          clones.s <- clones.s[which(clones.s[,2]==max(clones.s[,2],na.rm=TRUE)),]
      #if(nrow(clones.s)<1){print (paste("this:",sw,r,nmu,repl))}
          number.interfering.clones <- nrow(clones.s)
          ind <- match(clones.s[,1],dat[,1])
          avg.freq <- mean(dat[ind,2])
          sd.freq <- sd(dat[ind,2])
          column <- column+1
          res[row,column] <- number.interfering.clones
          column <- column+1
          res[row,column] <- avg.freq
          column <- column+1
          res[row,column] <- sd.freq
          column <- column+1
          # save maximum s attained at end of run, or 0
          if(nrow(clones.s)>0){
            res[row,column] <- clones.s[1,2]
          } else {
            res[row,column] <- 0
          }
        }
      }
      print (row)
    }   
  }
}

write.table(res,"ci.table.csv",sep=",",row.names=FALSE,col.names=FALSE,quote=FALSE)

# colors
cicols <- c(rgb(0,0,0),colorRampPalette( c("yellow", "red", "darkred"), space="rgb")(9))
# maximum s
max.s <- c((1+c(0.001,0.01,0.1,0.5,1))^5,(1+2)^4)

# three axes
dpl <- c(0,0.25,0.5,0.75,1)
scll <- c(0.1,0.5,1,2)
sml <- c(10^-5,10^-6,10^-7,10^-8)

pdf("clonal.interference.pdf",width=6,height=6)
s3d <- scatterplot3d(1,1,1,lab=c(4,5,4),box=FALSE,grid=TRUE,tick.marks=TRUE,
                     x.ticklabs=scll,y.ticklabs=dpl,z.ticklabs=sml,
                     xlab = "Selective coefficient", ylab="Crypt expansion probability", zlab="Selective mutation rate",
                     xlim=c(1,4),ylim=c(1,5),zlim=c(1,4),type="p",pch=NA,angle=55,col.axis="gray")
row <- 0
# z-axis is sel mu
for (smui in 1:length(sml)){
  smu <- sml[smui]
  # x-axis is sel. coeff.
  for (ss in 1:length(scl)){
    s <- scl[ss]
    # y-axis is r
    for (r in 1:length(dpl)){
      row <- row + 1
      if(s>=0.1){
        avg.cancer <- 0
      # color is average of 9 runs
        for(coli in 1:9){
          if(res[row,coli*4] %in% max.s){
            avg.cancer <- avg.cancer + 1
          }
        }
      #  print (avg.cancer)
      # draw perpendicular line to plane
        if(avg.cancer>0){
          s3d$points3d(rbind(c(ss-2,r,smui),c(ss-2,r,max(smui-1,1))), type = 'l', col="gray")
        }
                                        # x, y, z, pch, col
                                       #s3d$points3d(ss,r,smui,pch=16,col=cicols[avg.cancer])
        if(avg.cancer>0){
          s3d$points3d(ss-2,r,smui,pch=as.character(avg.cancer),col="black")
        } else {
          s3d$points3d(ss-2,r,smui,pch=NA,col=cicols[avg.cancer])
        }
      }
    }
  }
}
dev.off()






# extra

subres <- res[,1:9]
subres1 <- subres[,order(subres[1,]*1+subres[2,]*2+subres[3,]*4+subres[4,]*8+subres[5,]*16,decreasing=FALSE)]
subres <- res[,10:18]
subres2 <- subres[,order(subres[1,]*1+subres[2,]*2+subres[3,]*4+subres[4,]*8+subres[5,]*16,decreasing=FALSE)]
subres <- res[,19:27]
subres3 <- subres[,order(subres[1,]*1+subres[2,]*2+subres[3,]*4+subres[4,]*8+subres[5,]*16,decreasing=FALSE)]
subres <- res[,28:36]
subres4 <- subres[,order(subres[1,]*1+subres[2,]*2+subres[3,]*4+subres[4,]*8+subres[5,]*16,decreasing=FALSE)]
subres <- res[,37:45]
subres5 <- subres[,order(subres[1,]*1+subres[2,]*2+subres[3,]*4+subres[4,]*8+subres[5,]*16,decreasing=FALSE)]
nc <- t(t(c(NA,NA,NA,NA,NA,NA,NA,NA)))
ress <- cbind(subres1,nc,subres2,nc,subres3,nc,subres4,nc,subres5)
scheme <- c("magenta","blue","green","orange","red")
ress <- as.matrix(ress)
# split matrix, remove number of interfering clones
ci.labels <- ress[6:8,]
ress <- ress[1:5,]
rownames(ress) <- c("1 hit","2 hits","3 hits","4 hits", "5 hits")
pdf("beneficial.mutations.frequencies.sm10-7.s2.pdf",width=8, height=6)
mp <- barplot2(ress, beside = FALSE,col = scheme[1:5],ylim = c(0,1.1), plot.grid = FALSE,axes=FALSE)
axis(2,at=c(0,.25,.5,.75,1),label=c(0,.25,.5,.75,1))
axis(1,at=c(c(5,15,25,35,45)*1.2-0.5),label=c("r=0","r=.25","r=.5","r=.75","r=1"),tick=FALSE)
# calculate y-values
y.value.labels <- apply(ress[1:5,],2,sum)
text(x=c(1:49)*1.2-0.5,y=y.value.labels+0.05,ci.labels[1,1:45],cex=0.6)
#text(x=c(1:45)*1.2-0.5,y=y.value.labels+0.2,format(ci.labels[2,1:45],digits=1),cex=0.6)
#text(x=c(1:45)*1.2-0.5,y=y.value.labels+0.1,format(ci.labels[3,1:45],digits=1),cex=0.6)
#legend(24,1,legend=rownames(res),fill=gray(1-(0:4/4)))
mtext(side=2, text="Frequency of beneficial mutations",line=2)
mtext(side=1, text="Variation between individual runs",line=2)
dev.off()


